use ::alloc::alloc::{GlobalAlloc, Layout};
use core::{cmp::min, ptr::copy_nonoverlapping};
use ffi::*;

mod ffi {
    use crate::stdlib::lock::SpinLock;

    static mut LOCK: SpinLock = SpinLock::new();

    extern "C" {
        pub fn early_allocate_new_page(size: u32);
        pub fn kmalloc(size: usize) -> *mut u8;
        pub fn krealloc(pointer: *mut u8, size: usize) -> *mut u8;
        pub fn kcalloc(num: usize, size: usize) -> *mut u8;
        pub fn kfree(pointer: *mut u8);
    }

    #[no_mangle]
    pub extern "C" fn liballoc_lock() -> u32 {
        unsafe {
            LOCK.lock();
        }
        0
    }

    #[no_mangle]
    pub extern "C" fn liballoc_unlock() -> u32 {
        unsafe {
            LOCK.unlock();
        }
        0
    }

    #[no_mangle]
    pub extern "C" fn liballoc_alloc(num: usize) -> *mut u8 {
        die()
    }

    #[no_mangle]
    pub extern "C" fn liballoc_free(pointer: *mut u8, num: usize) -> u32 {
        die()
    }
}

pub struct Allocator;

const MIN_ALIGN: usize = 0x1;

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            kmalloc(layout.size())
        } else {
            alloc_aligned(layout)
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if layout.align() <= MIN_ALIGN {
            kfree(ptr)
        } else {
            free_aligned(ptr, layout)
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            krealloc(ptr, new_size)
        } else {
            let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
            let new_ptr = self.alloc(new_layout);
            if !new_ptr.is_null() {
                copy_nonoverlapping(ptr, new_ptr, min(layout.size(), new_size));
                self.dealloc(ptr, layout);
            }
            new_ptr
        }
    }
}

impl Allocator {
    pub fn init() {
        unsafe {
            early_allocate_new_page(0x20_0000);
        }
    }
}

unsafe fn alloc_aligned(layout: Layout) -> *mut u8 {
    let (size, align) = (layout.size(), layout.align());
    let total_size = size + (2 * align) + size_of::<*mut u8>();
    let ptr = kmalloc(total_size);

    if ptr.is_null() {
        ptr
    } else {
        let ptr_aligned = ((ptr as usize + align - 1) & (!(align - 1))) as *mut u8;
        let ptr_free = (ptr as usize + size) as *mut *mut u8;
        *ptr_free = ptr;
        ptr_aligned
    }
}

unsafe fn free_aligned(ptr: *mut u8, layout: Layout) {
    let size = layout.size();
    let ptr_free = (ptr as usize + size) as *mut *mut u8;
    let ptr = *ptr_free;
    kfree(ptr);
}

#[alloc_error_handler]
fn alloc_error_handler(layout: Layout) -> ! {
    die()
}
